'''
计算策略所需要的字段，但是在抓取中没有的字段，主要是两个字段pe_ttm和复权因子
computePeTtm:计算滚动市盈率，要注意，1.未来函数的问题， 2.财报不够四个季度的情况
computeAdjfctor:计算复权因子
'''

import sys

from pymongo import MongoClient, DESCENDING, UpdateOne

from pe_ttm.DataCrawler import getAllCodes

db = MongoClient('mongodb://localhost:27017')['quant']

'''
计算股票在某时的滚动市盈率
'''
def computePeTtm():
    # 获取所有股票
    codes = getAllCodes()
    codes = ['000720']
    for code in codes:
        print("计算滚动市盈率,{0}".format(code))

        dailyCursor = db.daily_k.find(
            {'code': code},
            projection={'close': True, 'time': True}
        )

        updateRequests = []
        for daily in dailyCursor:
            date = daily['time']

            financeReportCursor = db.finance_report.find(
                {'code': code, 'announced_date': {'$lte': date}},
                sort=[('announced_date', DESCENDING)],
                limit = 4
            )

            financeReports = [x for x in financeReportCursor]
            print(financeReports)

            # 如果数据不满4个季度，则不参与计算
            if len(financeReports) < 4:
                continue

            # 计算滚动市盈率并保存到daily_k中
            eps = 0
            for financeReport in financeReports:
                if financeReport['eps'] != '-':
                    eps += financeReport['eps']
            # 计算PE_TTM
            if eps != 0:
                updateRequests.append(
                    UpdateOne(
                        {'code': code, 'time': date},
                        {'$set':{'pe_ttm': round(daily['close']/eps, 4)}}
                    )
                )

        if len(updateRequests) > 0:
            updateResult = db.daily_k.bulk_write(updateRequests, ordered=False)
            print("更新PE_TTM,{0}， 更新:{1}".format(code, updateResult.modified_count))

'''
计算复权因子
'''
def computeAdjfactor():
    # 获取所有股票列表
    codes = getAllCodes()

    for code in codes:
        # 找到所有为复权的日K线
        dailyKCursor = db.daily_k.find(
            {'code': code},
            projection={'time': True, 'close': True},
            batch_size = 1000
        )

        # 创建time为key，daily_k为value的dict
        timeDailyKDict = dict()

        for dailyK in dailyKCursor:
            timeDailyKDict[dailyK['time']] = dailyK

        # 找到所有后复权的日K线
        dailyKBaCusor = db.daily_k_ba.find(
            {'code': code},
            projection = {'code': True, 'time': True, 'close': True},
            batch_size = 1000
        )

        # 按日期,计算复权因子
        updateRequest = []
        for dailyKBa in dailyKBaCusor:
            time = dailyKBa['time']
            if time in timeDailyKDict:
                dailyK = timeDailyKDict[time]
                # 后复权价格 = 未复权价格 * 复权因子
                adjfactor = dailyKBa['close'] / dailyK['close']

                updateRequest.append(
                    UpdateOne(
                        {'code': code, 'time': time},
                        {'$set': {'adjfactor': adjfactor}}
                    )
                )

        # 更新数据
        if len(updateRequest) > 0:
            result = db.daily_k.bulk_write(updateRequest, ordered=False)
            print("{0}, modified:{1}".format(code, result.modified_count))

if __name__ == '__main__':
    if len(sys.argv) <= 1:
        print("请指定类型adjfactor - 计算复权因子, pe - 计算滚动市盈率")
    else:
        type = sys.argv[1]

        if type == 'adjfactor':
            computeAdjfactor()
        elif type == 'pe':
            computePeTtm()
